import sys
assert sys.version_info.major == 3, 'ERROR: hj_dataset_devkit works on Python3 ONLY'
from enum import Enum, unique

import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from dataset_sdk import DatasetDevkit, Scene, Frame
from dataset_sdk import HongjingDevkit, HongjingScene, HongjingFrame
from dataset_sdk import NuscenesDevkit, NuscenesScene
try:
    from dataset_sdk import Argoverse2Devkit
except:
    Argoverse2Devkit = None
from dataset_sdk import CoordinateSystem
from dataset_sdk import ObstacleCategory, BoundingBox, Obstacle, MotionObstacle
from dataset_sdk import MapInfo
from dataset_sdk import LaneDivider, LaneDividerLineType, LaneDividerColor
from dataset_sdk import LaneDividerJoint
from dataset_sdk import RoadEdge, RoadEdgeCategory
from dataset_sdk import RoadMarking, RoadMarkingCategory
from dataset_sdk import TrafficSign, TrafficSignCategory
from dataset_sdk import TrafficLight, TrafficLightType, TrafficLightStatus
from dataset_sdk import CenterLine, CenterLineTopoPoint, CenterLineConnect

@unique
class SupportedDataset(Enum):
    hongjing=0
    nuscenes=1
    argoverse2=3

def load_dataset(dataset_type:SupportedDataset, path, **kwargs) -> DatasetDevkit:
    dataset_map = {
        SupportedDataset.hongjing: HongjingDevkit,
        SupportedDataset.nuscenes: NuscenesDevkit,
        SupportedDataset.argoverse2: Argoverse2Devkit,
    }
    return dataset_map[dataset_type](path, **kwargs)
